package main

import (
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"

	"gitee.com/yuebanlaosiji/gocelery/pkg/gocelery"
)

// Define tasks that can be called from Python
func multiplyTask(args []interface{}, kwargs map[string]interface{}) (interface{}, error) {
	if len(args) != 2 {
		return nil, fmt.Errorf("expected 2 arguments, got %d", len(args))
	}

	// Type assertions for the arguments
	a, ok := args[0].(float64)
	if !ok {
		return nil, fmt.Errorf("expected first argument to be a number")
	}

	b, ok := args[1].(float64)
	if !ok {
		return nil, fmt.Errorf("expected second argument to be a number")
	}

	result := a * b
	fmt.Printf("Multiply task executed: %f * %f = %f\n", a, b, result)
	return result, nil
}

func greetTask(args []interface{}, kwargs map[string]interface{}) (interface{}, error) {
	var name string

	// Check if name was passed as a positional argument
	if len(args) > 0 {
		if n, ok := args[0].(string); ok {
			name = n
		}
	}

	// Check if name was passed as a keyword argument
	if n, ok := kwargs["name"].(string); ok {
		name = n
	}

	if name == "" {
		name = "Guest"
	}

	greeting := fmt.Sprintf("Hello, %s from Go!", name)
	fmt.Println(greeting)

	return greeting, nil
}

func addTask(args []interface{}, kwargs map[string]interface{}) (interface{}, error) {
	if len(args) != 2 {
		return nil, fmt.Errorf("expected 2 arguments, got %d", len(args))
	}

	result := args[0].(float64) + args[1].(float64)
	return result, nil
}

func main() {
	// Determine if we should run as a worker or client
	// Usage: ./python_interop worker|client [task_name]
	if len(os.Args) < 2 {
		fmt.Println("Usage: ./python_interop worker|client [task_name]")
		os.Exit(1)
	}

	mode := os.Args[1]

	// Configure Redis connection URL
	// For Redis without password:
	// redisURL := "redis://localhost:6379/0"

	// For Redis with password:
	redisURL := "redis://:Hewenyu1994@121.5.249.100:6379/0"

	// For Redis with username and password (Redis 6.0+):
	// redisURL := "redis://username:password@localhost:6379/0"

	// Configuration matching typical Python Celery defaults
	config := gocelery.Config{
		BrokerURL:    redisURL,
		BackendURL:   redisURL,
		DefaultQueue: "celery",
		Concurrency:  2,
	}

	switch mode {
	case "worker":
		runWorker(config)
	case "client":
		if len(os.Args) < 3 {
			fmt.Println("Usage: ./python_interop client <task_name>")
			fmt.Println("Available tasks: multiply, greet, python.add, python.process")
			os.Exit(1)
		}
		taskName := os.Args[2]
		runClient(config, taskName)
	default:
		fmt.Printf("Unknown mode: %s. Use 'worker' or 'client'\n", mode)
		os.Exit(1)
	}
}

func runWorker(config gocelery.Config) {
	// Create the worker
	worker, err := gocelery.NewWorker(config)
	if err != nil {
		log.Fatalf("Failed to create worker: %v", err)
	}

	// Register task handlers
	err = worker.Register("tasks.multiply", multiplyTask)
	if err != nil {
		log.Fatalf("Failed to register multiply task: %v", err)
	}

	err = worker.Register("tasks.greet", greetTask)
	if err != nil {
		log.Fatalf("Failed to register greet task: %v", err)
	}

	err = worker.Register("tasks.add", addTask)
	if err != nil {
		log.Fatalf("Failed to register add task: %v", err)
	}

	fmt.Println("Starting worker...")
	err = worker.Start()
	if err != nil {
		log.Fatalf("Failed to start worker: %v", err)
	}

	// Handle graceful shutdown
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	fmt.Println("Worker started. Press Ctrl+C to stop.")
	fmt.Println("Registered tasks:")
	fmt.Println("- tasks.multiply")
	fmt.Println("- tasks.greet")
	fmt.Println("- tasks.add")

	<-sigChan
	fmt.Println("\nShutting down...")
	worker.Stop()
	fmt.Println("Worker stopped.")
}

func runClient(config gocelery.Config, taskName string) {
	// Create the client
	client, err := gocelery.NewClient(config)
	if err != nil {
		log.Fatalf("Failed to create client: %v", err)
	}
	defer client.Close()

	var taskID string

	// Call the appropriate task
	switch taskName {
	case "multiply":
		taskID, err = client.SendTask("tasks.multiply", []interface{}{5, 7}, nil)
	case "greet":
		taskID, err = client.SendTask("tasks.greet", []interface{}{"GoCelery"}, nil)
	case "python.add":
		// 使用正确的Python Celery任务名称格式
		taskID, err = client.SendTask("tasks.add", []interface{}{10, 20}, nil)
	case "python.process":
		// 使用通用处理器
		taskID, err = client.SendTask("tasks.process_any", []interface{}{10, 20}, map[string]interface{}{
			"task_name": "add",
		})
	default:
		log.Fatalf("Unknown task: %s", taskName)
	}

	if err != nil {
		log.Fatalf("Failed to send task: %v", err)
	}

	fmt.Printf("Task sent. ID: %s\n", taskID)

	// Wait for the result
	fmt.Println("Waiting for result...")
	result, err := client.WaitForResult(taskID, 10*time.Second)
	if err != nil {
		log.Fatalf("Failed to get result: %v", err)
	}

	fmt.Printf("Task result: %v\n", result)
}
